#include <rclcpp/rclcpp.hpp>
#include <nav_msgs/msg/odometry.hpp>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <std_msgs/msg/bool.hpp>
#include <Eigen/Dense>
#include <vector>
#include <fstream>
#include <string>
#include <random>
#include <algorithm>
#include <cmath>
#include <numeric>
#include <message_filters/subscriber.h>
#include <message_filters/sync_policies/approximate_time.h>
#include <message_filters/synchronizer.h>

class OnlineOdometryCalibrator : public rclcpp::Node
{
public:
    OnlineOdometryCalibrator()
        : Node("online_odometry_calibrator"),
          mode_(Mode::ONLINE)
    {
        gps_sub_.subscribe(this, "/odometry/gps");
        lo_sub_.subscribe(this, "/Odometry/transform");

        sync_ = std::make_shared<message_filters::Synchronizer<MySyncPolicy>>(
            MySyncPolicy(100), gps_sub_, lo_sub_);
        sync_->registerCallback(std::bind(&OnlineOdometryCalibrator::callback, this, std::placeholders::_1, std::placeholders::_2));

        extrinsic_pub_ = this->create_publisher<geometry_msgs::msg::TransformStamped>("extrinsics", 10);

        control_sub_ = this->create_subscription<std_msgs::msg::Bool>(
            "align_all", 10,
            std::bind(&OnlineOdometryCalibrator::controlCallback, this, std::placeholders::_1));

        window_size_ = this->declare_parameter<int>("window_size", 80);
        min_points_ = this->declare_parameter<int>("min_points", 10);
        cache_file_ = this->declare_parameter<std::string>("cache_file", "odom_data_cache.txt");
        history_window_size_ = this->declare_parameter<int>("history_window_size", 300);

        max_gps_jump_ = this->declare_parameter<double>("max_gps_jump", 5.0);       // meters
        min_moving_dist_ = this->declare_parameter<double>("min_moving_dist", 0.1); // meters
        ransac_iters_ = this->declare_parameter<int>("ransac_iters", 100);
        ransac_thresh_ = this->declare_parameter<double>("ransac_thresh", 2.5); // meters

        time_weight_lambda_ = this->declare_parameter<double>("time_weight_lambda", 0.0007); // 权重衰减系数

        RCLCPP_INFO(this->get_logger(), "在线里程计标定节点已启动。高级滤波（加权最小二乘）历史标定模式已启用。");
    }

private:
    enum class Mode
    {
        ONLINE,
        HISTORY
    };
    Mode mode_;

    typedef message_filters::sync_policies::ApproximateTime<
        nav_msgs::msg::Odometry, nav_msgs::msg::Odometry>
        MySyncPolicy;

    message_filters::Subscriber<nav_msgs::msg::Odometry> gps_sub_;
    message_filters::Subscriber<nav_msgs::msg::Odometry> lo_sub_;
    std::shared_ptr<message_filters::Synchronizer<MySyncPolicy>> sync_;

    std::vector<Eigen::Vector3d> gps_points_;
    std::vector<Eigen::Vector3d> lo_points_;
    std::vector<Eigen::Vector3d> all_gps_points_;
    std::vector<Eigen::Vector3d> all_lo_points_;
    std::vector<rclcpp::Time> all_times_;

    int window_size_;
    int min_points_;
    int history_window_size_;
    std::string cache_file_;

    double max_gps_jump_;
    double min_moving_dist_;
    int ransac_iters_;
    double ransac_thresh_;
    double time_weight_lambda_;

    Eigen::Matrix3d fixed_R_;
    Eigen::Vector3d fixed_t_;
    bool has_history_extrinsic_ = false;

    rclcpp::Publisher<geometry_msgs::msg::TransformStamped>::SharedPtr extrinsic_pub_;
    rclcpp::Subscription<std_msgs::msg::Bool>::SharedPtr control_sub_;

    void controlCallback(const std_msgs::msg::Bool::SharedPtr msg)
    {
        if (msg->data)
        {
            if (mode_ != Mode::HISTORY)
            {
                RCLCPP_INFO(this->get_logger(), "切换到历史标定模式，使用高级滤波和加权鲁棒配准。");
                mode_ = Mode::HISTORY;

                std::vector<Eigen::Vector3d> cleaned_gps, cleaned_lo;
                std::vector<rclcpp::Time> cleaned_times;
                cleanHistoryData(cleaned_lo, cleaned_gps, cleaned_times);

                if (cleaned_gps.size() > size_t(history_window_size_))
                {
                    cleaned_gps = std::vector<Eigen::Vector3d>(cleaned_gps.end() - history_window_size_, cleaned_gps.end());
                    cleaned_lo = std::vector<Eigen::Vector3d>(cleaned_lo.end() - history_window_size_, cleaned_lo.end());
                    cleaned_times = std::vector<rclcpp::Time>(cleaned_times.end() - history_window_size_, cleaned_times.end());
                }

                if (cleaned_gps.size() >= size_t(min_points_))
                {
                    // Step 1: RANSAC剔除异常点
                    Eigen::Matrix3d ransac_R;
                    Eigen::Vector3d ransac_t;
                    std::vector<int> inlier_indices;
                    computeUmeyamaRANSACWithInliers(cleaned_lo, cleaned_gps, ransac_R, ransac_t, ransac_iters_, ransac_thresh_, inlier_indices);

                    // Step 2: 对RANSAC内点加权最小二乘滤波优化（时间加权）
                    std::vector<Eigen::Vector3d> inlier_lo, inlier_gps;
                    std::vector<rclcpp::Time> inlier_times;
                    for (int idx : inlier_indices)
                    {
                        inlier_lo.push_back(cleaned_lo[idx]);
                        inlier_gps.push_back(cleaned_gps[idx]);
                        if (idx < int(cleaned_times.size()))
                            inlier_times.push_back(cleaned_times[idx]);
                    }

                    computeWeightedUmeyama(inlier_lo, inlier_gps, inlier_times, fixed_R_, fixed_t_, time_weight_lambda_);
                    has_history_extrinsic_ = true;
                    publishExtrinsic(fixed_R_, fixed_t_, now());
                    printExtrinsic(fixed_R_, fixed_t_);
                    saveCache(cache_file_, inlier_lo, inlier_gps);
                }
                else
                {
                    RCLCPP_WARN(this->get_logger(), "高质量历史数据点不足，无法进行历史配准。");
                    has_history_extrinsic_ = false;
                }
            }
        }
        else
        {
            if (mode_ != Mode::ONLINE)
            {
                RCLCPP_INFO(this->get_logger(), "切换到在线标定模式，实时输出滑动窗口外参。");
                mode_ = Mode::ONLINE;
                has_history_extrinsic_ = false;
                if (gps_points_.size() >= size_t(min_points_))
                {
                    Eigen::Matrix3d R;
                    Eigen::Vector3d t;
                    computeUmeyama(lo_points_, gps_points_, R, t);
                    publishExtrinsic(R, t, now());
                    printExtrinsic(R, t);
                }
            }
        }
    }

    void callback(const nav_msgs::msg::Odometry::ConstSharedPtr gps_msg,
                  const nav_msgs::msg::Odometry::ConstSharedPtr lo_msg)
    {
        Eigen::Vector3d gps_pt(gps_msg->pose.pose.position.x,
                               gps_msg->pose.pose.position.y,
                               gps_msg->pose.pose.position.z);
        Eigen::Vector3d lo_pt(lo_msg->pose.pose.position.x,
                              lo_msg->pose.pose.position.y,
                              lo_msg->pose.pose.position.z);

        gps_points_.push_back(gps_pt);
        lo_points_.push_back(lo_pt);

        if (gps_points_.size() > size_t(window_size_))
        {
            gps_points_.erase(gps_points_.begin());
            lo_points_.erase(lo_points_.begin());
        }

        all_gps_points_.push_back(gps_pt);
        all_lo_points_.push_back(lo_pt);
        all_times_.push_back(gps_msg->header.stamp);

        if (mode_ == Mode::ONLINE && gps_points_.size() >= size_t(min_points_))
        {
            Eigen::Matrix3d R;
            Eigen::Vector3d t;
            computeUmeyama(lo_points_, gps_points_, R, t);
            publishExtrinsic(R, t, lo_msg->header.stamp);
            printExtrinsic(R, t);
        }
    }

    void cleanHistoryData(std::vector<Eigen::Vector3d> &clean_lo, std::vector<Eigen::Vector3d> &clean_gps, std::vector<rclcpp::Time> &clean_times)
    {
        clean_lo.clear();
        clean_gps.clear();
        clean_times.clear();
        size_t N = all_gps_points_.size();
        if (N < 2)
            return;
        clean_gps.push_back(all_gps_points_[0]);
        clean_lo.push_back(all_lo_points_[0]);
        clean_times.push_back(all_times_[0]);
        for (size_t i = 1; i < N; ++i)
        {
            double gps_jump = (all_gps_points_[i] - all_gps_points_[i - 1]).norm();
            double lo_jump = (all_lo_points_[i] - all_lo_points_[i - 1]).norm();
            if (gps_jump < max_gps_jump_ && gps_jump > min_moving_dist_ && lo_jump > min_moving_dist_)
            {
                clean_gps.push_back(all_gps_points_[i]);
                clean_lo.push_back(all_lo_points_[i]);
                clean_times.push_back(all_times_[i]);
            }
        }
    }

    void computeUmeyama(const std::vector<Eigen::Vector3d> &src, const std::vector<Eigen::Vector3d> &dst, Eigen::Matrix3d &R, Eigen::Vector3d &t)
    {
        size_t N = src.size();
        Eigen::MatrixXd src_mat(3, N), dst_mat(3, N);
        for (size_t i = 0; i < N; ++i)
        {
            src_mat.col(i) = src[i];
            dst_mat.col(i) = dst[i];
        }
        Eigen::Vector3d src_mean = src_mat.rowwise().mean();
        Eigen::Vector3d dst_mean = dst_mat.rowwise().mean();
        src_mat.colwise() -= src_mean;
        dst_mat.colwise() -= dst_mean;

        Eigen::Matrix3d H = src_mat * dst_mat.transpose();
        Eigen::JacobiSVD<Eigen::MatrixXd> svd(H, Eigen::ComputeFullU | Eigen::ComputeFullV);
        Eigen::Matrix3d U = svd.matrixU();
        Eigen::Matrix3d V = svd.matrixV();
        R = V * U.transpose();
        if (R.determinant() < 0)
        {
            V.col(2) *= -1;
            R = V * U.transpose();
        }
        t = dst_mean - R * src_mean;
    }

    // RANSAC-Umeyama配准并输出内点索引
    void computeUmeyamaRANSACWithInliers(const std::vector<Eigen::Vector3d> &src,
                                         const std::vector<Eigen::Vector3d> &dst,
                                         Eigen::Matrix3d &best_R, Eigen::Vector3d &best_t,
                                         int max_iters, double thresh,
                                         std::vector<int> &best_inlier_indices)
    {
        if (src.size() < 3 || dst.size() < 3)
        {
            RCLCPP_WARN(this->get_logger(), "RANSAC数据点不足");
            best_R = Eigen::Matrix3d::Identity();
            best_t = Eigen::Vector3d::Zero();
            best_inlier_indices.clear();
            return;
        }
        std::vector<int> idx(src.size());
        std::iota(idx.begin(), idx.end(), 0);
        std::random_device rd;
        std::mt19937 g(rd());

        int best_inliers = 0;
        double best_error = 1e8;

        for (int iter = 0; iter < max_iters; ++iter)
        {
            std::shuffle(idx.begin(), idx.end(), g);
            std::vector<Eigen::Vector3d> src_sub, dst_sub;
            for (int k = 0; k < 3; ++k)
            {
                src_sub.push_back(src[idx[k]]);
                dst_sub.push_back(dst[idx[k]]);
            }
            Eigen::Matrix3d R;
            Eigen::Vector3d t;
            computeUmeyama(src_sub, dst_sub, R, t);

            int inliers = 0;
            double total_err = 0.0;
            std::vector<int> inlier_indices;
            for (size_t i = 0; i < src.size(); ++i)
            {
                Eigen::Vector3d pred = R * src[i] + t;
                double err = (pred - dst[i]).norm();
                if (err < thresh)
                {
                    inliers++;
                    inlier_indices.push_back(i);
                    total_err += err;
                }
            }
            if (inliers > best_inliers || (inliers == best_inliers && total_err < best_error))
            {
                if (inliers >= 3)
                {
                    std::vector<Eigen::Vector3d> src_inlier, dst_inlier;
                    for (int ix : inlier_indices)
                    {
                        src_inlier.push_back(src[ix]);
                        dst_inlier.push_back(dst[ix]);
                    }
                    computeUmeyama(src_inlier, dst_inlier, best_R, best_t);
                    best_inliers = inliers;
                    best_error = total_err;
                    best_inlier_indices = inlier_indices;
                }
            }
        }
        RCLCPP_INFO(this->get_logger(), "RANSAC历史标定，最佳内点数=%d，总误差=%.4f", best_inliers, best_error);
    }

    // 加权Umeyama（适合滤波优化）
    void computeWeightedUmeyama(const std::vector<Eigen::Vector3d> &src,
                                const std::vector<Eigen::Vector3d> &dst,
                                const std::vector<rclcpp::Time> &times,
                                Eigen::Matrix3d &R, Eigen::Vector3d &t,
                                double lambda)
    {
        size_t N = src.size();
        if (N < 3 || dst.size() < 3 || times.size() < 3)
        {
            computeUmeyama(src, dst, R, t);
            return;
        }
        std::vector<double> weights(N, 1.0);

        // 以最新时刻权重最大，历史越久权重越小
        rclcpp::Time newest = times.back();
        for (size_t i = 0; i < N; ++i)
        {
            double dt = (newest - times[i]).seconds();
            weights[i] = std::exp(-lambda * dt);
        }
        double sum_w = std::accumulate(weights.begin(), weights.end(), 0.0);
        for (size_t i = 0; i < N; ++i)
            weights[i] /= sum_w;

        Eigen::Vector3d src_mean = Eigen::Vector3d::Zero();
        Eigen::Vector3d dst_mean = Eigen::Vector3d::Zero();
        for (size_t i = 0; i < N; ++i)
        {
            src_mean += weights[i] * src[i];
            dst_mean += weights[i] * dst[i];
        }
        Eigen::MatrixXd src_mat(3, N), dst_mat(3, N);
        for (size_t i = 0; i < N; ++i)
        {
            src_mat.col(i) = src[i] - src_mean;
            dst_mat.col(i) = dst[i] - dst_mean;
        }
        for (size_t i = 0; i < N; ++i)
        {
            src_mat.col(i) *= sqrt(weights[i]);
            dst_mat.col(i) *= sqrt(weights[i]);
        }
        Eigen::Matrix3d H = src_mat * dst_mat.transpose();
        Eigen::JacobiSVD<Eigen::MatrixXd> svd(H, Eigen::ComputeFullU | Eigen::ComputeFullV);
        Eigen::Matrix3d U = svd.matrixU();
        Eigen::Matrix3d V = svd.matrixV();
        R = V * U.transpose();
        if (R.determinant() < 0)
        {
            V.col(2) *= -1;
            R = V * U.transpose();
        }
        t = dst_mean - R * src_mean;
    }

    void publishExtrinsic(const Eigen::Matrix3d &R, const Eigen::Vector3d &t, const rclcpp::Time &stamp)
    {
        geometry_msgs::msg::TransformStamped msg;
        msg.header.stamp = stamp;
        msg.header.frame_id = "lo_frame";
        msg.child_frame_id = "gps_frame";
        msg.transform.translation.x = t.x();
        msg.transform.translation.y = t.y();
        msg.transform.translation.z = t.z();
        Eigen::Quaterniond q(R);
        msg.transform.rotation.x = q.x();
        msg.transform.rotation.y = q.y();
        msg.transform.rotation.z = q.z();
        msg.transform.rotation.w = q.w();
        extrinsic_pub_->publish(msg);
    }

    void printExtrinsic(const Eigen::Matrix3d &R, const Eigen::Vector3d &t)
    {
        RCLCPP_INFO(this->get_logger(), "当前外参旋转矩阵R：\n%.6f %.6f %.6f\n%.6f %.6f %.6f\n%.6f %.6f %.6f",
                    R(0, 0), R(0, 1), R(0, 2),
                    R(1, 0), R(1, 1), R(1, 2),
                    R(2, 0), R(2, 1), R(2, 2));
        RCLCPP_INFO(this->get_logger(), "当前外参平移向量t： [%.6f, %.6f, %.6f]",
                    t.x(), t.y(), t.z());
    }

    void saveCache(const std::string &filename,
                   const std::vector<Eigen::Vector3d> &lo_points,
                   const std::vector<Eigen::Vector3d> &gps_points)
    {
        std::ofstream ofs(filename);
        if (!ofs.is_open())
        {
            RCLCPP_WARN(this->get_logger(), "无法打开缓存文件进行写入：%s", filename.c_str());
            return;
        }
        ofs << "# lo_x lo_y lo_z gps_x gps_y gps_z\n";
        for (size_t i = 0; i < gps_points.size(); ++i)
        {
            ofs << lo_points[i].x() << " " << lo_points[i].y() << " " << lo_points[i].z() << " "
                << gps_points[i].x() << " " << gps_points[i].y() << " " << gps_points[i].z() << "\n";
        }
        ofs.close();
        RCLCPP_INFO(this->get_logger(), "已保存 %zu 组对齐数据到文件：%s", gps_points.size(), filename.c_str());
    }
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<OnlineOdometryCalibrator>());
    rclcpp::shutdown();
    return 0;
}
